# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""
import pandas as pd
from tqdm import tqdm
from datetime import datetime, timedelta
from WindPy import w

now = datetime.now()

if not  w.isconnected():
    w.start()


def get_kline_data_by_wind_code(wind_code, mode):
    tm = str(now + timedelta(days=1))[:10]
    wsd_data = w.wsd(wind_code,"open,high,low,close,volume",
                     "2020-01-01",tm,mode)
    
    if wsd_data.ErrorCode == 0:
        df = pd.DataFrame(wsd_data.Data, 
                          index=wsd_data.Fields, 
                          columns=wsd_data.Times).T
        df.columns = [s.lower() for s in df.columns]
    else:
        df = pd.DataFrame()
    return df


def filter_stocks():
    stock_set = w.wset("sectorconstituent","date=2024-09-20;sectorid=a001010100000000")
    all_ipo_dates = {}
    if stock_set.ErrorCode ==0:
        stock_name_df = pd.DataFrame(stock_set.Data, columns=stock_set.Codes,
                                     index=stock_set.Fields).T
        all_chinese_names = dict(zip(stock_name_df['wind_code'], 
                                     stock_name_df['sec_name']))
        codes = stock_name_df['wind_code'].to_list()
        sub_codes_li = [codes[i:i+400] for i in range(0, len(codes), 400)]
        for sub_codes in tqdm(sub_codes_li):
            wind_code_string = ','.join(sub_codes)
            ipo = w.wss(wind_code_string, "ipo_date")
            ipo_dates = dict(zip(ipo.Codes, ipo.Data[0]))
            all_ipo_dates = {**all_ipo_dates, **ipo_dates}
        return all_ipo_dates, all_chinese_names


def process_k_data(kdata, code, name):
    # 添加均线
    
    ma_list = [5, 10, 20, 30, 60]
    for ma in ma_list:
        kdata[f"MA{ma}"] = kdata['close'].rolling(ma).mean()
    
    # 索引近期60日股票数据
    kdata = kdata.iloc[60:,:]
    df = pd.DataFrame(index=kdata.index)
    df['code'] = code
    df['name'] = name
    
    # 查看均线是否向上
    for ma in ma_list:
        df[f'c>{ma}'] = kdata['close'] > kdata[f'MA{ma}']
    for ma in ma_list:
        df[f"{ma}↑"] = kdata[f'MA{ma}'].diff(1) > 0
    
    # 查看过去5日均线是否向上或向下
    for ma in [5]:
        for shift_i in range(1, 5):
            df[f"MA{ma}lag{shift_i}↑"] = kdata[f'MA{ma}'].diff(1).shift(shift_i) > 0
    # 查看过去20日均线是否向上或向下
    for ma in [20]:
        for shift_i in range(1, 5):
            df[f"MA{ma}lag{shift_i}↑"] = kdata[f'MA{ma}'].diff(1).shift(shift_i) > 0        
    
    # 5日均线上传10,20,30,60日均线
    for ma in [10, 20, 30, 60]:
        kdata[f"MA5-MA{ma}"] = kdata['MA5'] - kdata[f'MA{ma}']
        kdata[f"MA5-MA{ma}_up"] = kdata[f"MA5-MA{ma}"] > 0
        kdata.loc[:,f"MA5-MA{ma}_shift1_up"] = kdata[f"MA5-MA{ma}"].shift(1) > 0
    
    kdata['volume60'] = kdata['volume'].rolling(60).mean()
    
    df = df.iloc[60:,:]
    kdata = kdata.iloc[60:,:]
    
    for ma in [10, 20, 30, 60]:
        df[f'5C{ma}'] = kdata[f"MA5-MA{ma}_up"] & kdata[f"MA5-MA{ma}_shift1_up"]
    df['主力介入'] = kdata['volume'] > kdata['volume60'] * 2 
    df['date'] = df.index
    return df, kdata


if __name__ == "__main__":
    mode = "Period=W;PriceAdj=F" # D是日度,W是周度
    ipo_dates, chinese_names = filter_stocks()
    collector = []
    for i, (wind_code, ipo_date) in enumerate(ipo_dates.items()):
        if ipo_date.year < 2020:
            print(i,wind_code, ipo_date)
            kline = get_kline_data_by_wind_code(wind_code,mode)
            signal, kline = process_k_data(kline, wind_code, 
                                           chinese_names[wind_code])
            collector.append(signal)
    ot_df = pd.concat(collector, axis=0)
    ot_df['date'] = pd.to_datetime(ot_df['date'])
    ot_df = ot_df[ot_df['date'] > pd.to_datetime('2024-01-01')]
    ot_df.to_excel('stock_W-241008.xlsx')